Create VJP perf fix #4806

mlxd · 2023-11-08T15:56:31Z

Before submitting

Please complete the following checklist when submitting a PR:

All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to the
test directory!
All new functions and code must be clearly commented and documented.
If you do make documentation changes, make sure that the docs build and
render correctly by running make docs.
Ensure that the test suite passes, by running make test.
Add a new entry to the doc/releases/changelog-dev.md file, summarizing the
change, and including a link back to the PR.
The PennyLane source code conforms to
PEP8 standards.
We check all of our code against Pylint.
To lint modified files, simply pip install pylint, and then
run pylint pennylane/path/to/file.py.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

Context: v0.33.1 backport of PR #4792

Description of the Change: Updates VJP pipeline to favour direct matrix-vector products where possible.

Benefits: Improves performance for many parameter/many observable workloads

Possible Drawbacks:

Related GitHub Issues: #4792

github-actions · 2023-11-08T15:56:49Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

trbromley

Thanks @mlxd @albi3ro! Nothing blocking from my side but I do have a few questions.

pennylane/gradients/vjp.py

codecov · 2023-11-08T18:09:16Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (1b20871) 99.64% compared to head (4136300) 99.64%.
Report is 2 commits behind head on v0.33.1-rc0.

Additional details and impacted files

@@             Coverage Diff              @@
##           v0.33.1-rc0    #4806   +/-   ##
============================================
  Coverage        99.64%   99.64%           
============================================
  Files              380      380           
  Lines            34266    34285   +19     
============================================
+ Hits             34145    34164   +19     
  Misses             121      121

Files	Coverage Δ
pennylane/_version.py	`100.00% <100.00%> (ø)`
pennylane/devices/qubit/simulate.py	`100.00% <100.00%> (ø)`
pennylane/gradients/vjp.py	`100.00% <100.00%> (ø)`
pennylane/math/multi_dispatch.py	`99.59% <100.00%> (+<0.01%)`	⬆️
pennylane/ops/op_math/symbolicop.py	`100.00% <100.00%> (ø)`
pennylane/transforms/defer_measurements.py	`100.00% <100.00%> (ø)`

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

rmoyard

Thanks for the changes! It would be nice to have the list of cases where the it becomes slower because of the fallback. What I can identify is the following:

The performance is better:

Single measurement (not tensorflow) with parameters of the same shape.
Multiple measurements with same shapes and with parameters of the same shape.

The performance is worse:

Parameters of the different shapes.
Tensorflow in general takes a hit because it does not support @ for mulitplication
Mixing measurements with different shapes is slower because of the fallback in the compute_vjp_multi.

In the future we should come back to dispatch the cases properly.

pennylane/gradients/vjp.py

timmysilv · 2023-11-08T19:37:17Z

try-except generally hurts performance. I made a table of when we need which try-excepts so we can improve this in the future:

	non-TF	TF
non-ragged	none	both
ragged	multi	both

rmoyard

Thanks! We made a small table of what cases need what try expect and it also show which one are better and which one become worse. It will be easier to come back to this issue.

albi3ro

Performance improvements for the vast majority of hpc cases we care about, at the cost of a penality for edge cases and tensorflow. I find that to be an acceptable tradeoff.

While this code could probably always use more fine-tuning, this makes accomplishes exactly what it needs to accomplish,

Note that these changes have been reviewed and released already, they just aren't on master. **Context:** Bugfix release v0.33.1 had a change that isn't on master, so this gets them onto master **Description of the Change:** All I did was run `git merge v0.33.1`, then resolve conflicts. **Benefits:** Docs are in order, and the performance patch is on master. **Possible Drawbacks:** N/A For reference, these are the changes that were added in v0.33.1: v0.33.0...v0.33.1 - this PR has less files because all of the little fixes were already on master. It was only #4806 that never made it to master. --------- Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: Romain Moyard <rmoyard@gmail.com> Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> Co-authored-by: Isaac De Vlugt <isaacdevlugt@gmail.com> Co-authored-by: Isaac De Vlugt <34751083+isaacdevlugt@users.noreply.github.com> Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> Co-authored-by: Lee James O'Riordan <mlxd@users.noreply.github.com> Co-authored-by: BM7878 <117289949+BM7878@users.noreply.github.com> Co-authored-by: Josh Izaac <josh146@gmail.com>

Create VJP perf fix

36dce77

mlxd added gpu ci:run-full-test-suite Run the full test-suite on the pull request labels Nov 8, 2023

mlxd requested review from rmoyard, albi3ro and timmysilv November 8, 2023 15:58

trbromley reviewed Nov 8, 2023

View reviewed changes

pennylane/gradients/vjp.py Show resolved Hide resolved

pennylane/gradients/vjp.py Show resolved Hide resolved

pennylane/gradients/vjp.py Outdated Show resolved Hide resolved

timmysilv added 2 commits November 8, 2023 12:01

fix pylint disable with python 3.9

dfabb06

Merge branch 'v0.33.1-rc0' into fix/vjp_perf_ref_v0.33.1

347df6e

also run data tests for 3.9 in full test suite

4136300

rmoyard reviewed Nov 8, 2023

View reviewed changes

pennylane/gradients/vjp.py Show resolved Hide resolved

pennylane/gradients/vjp.py Show resolved Hide resolved

pennylane/gradients/vjp.py Show resolved Hide resolved

rmoyard self-requested a review November 8, 2023 19:41

rmoyard approved these changes Nov 8, 2023

View reviewed changes

albi3ro approved these changes Nov 8, 2023

View reviewed changes

mlxd added 2 commits November 8, 2023 15:08

Update changelog

fa97d53

[skip-ci]

96070f2

mlxd merged commit f3b41b6 into v0.33.1-rc0 Nov 8, 2023
33 of 37 checks passed

mlxd deleted the fix/vjp_perf_ref_v0.33.1 branch November 8, 2023 20:10

timmysilv mentioned this pull request Nov 14, 2023

Merge 0.33.1 fixes into master #4838

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Create VJP perf fix #4806

Create VJP perf fix #4806

mlxd commented Nov 8, 2023 •

edited

Loading

github-actions bot commented Nov 8, 2023

trbromley left a comment

codecov bot commented Nov 8, 2023 •

edited

Loading

rmoyard left a comment •

edited

Loading

timmysilv commented Nov 8, 2023

rmoyard left a comment

albi3ro left a comment

Create VJP perf fix #4806

Create VJP perf fix #4806

Conversation

mlxd commented Nov 8, 2023 • edited Loading

Before submitting

github-actions bot commented Nov 8, 2023

trbromley left a comment

Choose a reason for hiding this comment

codecov bot commented Nov 8, 2023 • edited Loading

Codecov Report

rmoyard left a comment • edited Loading

Choose a reason for hiding this comment

timmysilv commented Nov 8, 2023

rmoyard left a comment

Choose a reason for hiding this comment

albi3ro left a comment

Choose a reason for hiding this comment

mlxd commented Nov 8, 2023 •

edited

Loading

codecov bot commented Nov 8, 2023 •

edited

Loading

rmoyard left a comment •

edited

Loading